import torch
import torch.nn
import torch.nn.functional
import torch.optim
import torchvision
import torch.utils.tensorboard
import sys


# 定义网络模型, 继承 torch.nn.Module
class TestNet(torch.nn.Module):
    # 构造方法
    def __init__(self):
        # 调用父类构造方法, 也可以使用super关键字
        super().__init__()
        # 定义 Conv2d 卷积层
        self.conv1 = torch.nn.Conv2d(3, 6, 5)
        self.pool = torch.nn.MaxPool2d(2, 2)
        self.conv2 = torch.nn.Conv2d(6, 16, 5)
        self.fc1 = torch.nn.Linear(16 * 5 * 5, 120)
        self.fc2 = torch.nn.Linear(120, 84)
        self.fc3 = torch.nn.Linear(84, 10)

    def forward(self, x):
        # 卷积后跟一个 ReLU 函数
        x = self.pool(torch.nn.functional.relu(self.conv1(x)))
        x = self.pool(torch.nn.functional.relu(self.conv2(x)))
        # view是 numpy 的reshape的 PyTorch 版本
        x = x.view(-1, 16 * 5 * 5)
        x = torch.nn.functional.relu(self.fc1(x))
        x = torch.nn.functional.relu(self.fc2(x))
        x = self.fc3(x)
        return x


if __name__ == '__main__':

    # GPU/CPU?
    print(torch.cuda.is_available())
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print(device)

    # 创建网络实例
    testNet = TestNet()
    # 将模型参数发送给GPU
    testNet.to(device)
    # 定义损失
    lossFunction = torch.nn.CrossEntropyLoss()
    # 定义优化器
    optimizer = torch.optim.Adam(testNet.parameters())

    # 数据
    transform = torchvision.transforms.Compose(
        [torchvision.transforms.ToTensor(),
         torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
    trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=True, num_workers=2)

    # 可视化
    writer = torch.utils.tensorboard.SummaryWriter("./tensorBoard")
    # dataiter = iter(trainloader)
    # images, labels = dataiter.next()
    # writer.add_graph(testNet, images)

    # 训练
    for epoch in range(1):
        running_loss = 0.0
        for i, data in enumerate(trainloader, 0):
            inputs, labels = data
            # 将输入和输入也发送给GPU, 将张量(数组)转化为cuda数据类型
            inputs, labels = inputs.to(device), labels.to(device)
            # 清空损失, 否则会累积
            optimizer.zero_grad()
            # 正向传播
            outputs = testNet(inputs)
            # 反向传播
            loss = lossFunction(outputs, labels)
            loss.backward()
            # 优化, 更新参数权重
            optimizer.step()
            # 训练过程可视化
            running_loss += loss.item()
            if i % 2000 == 1999:
                print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 2000))
                running_loss = 0.0
            writer.add_scalar('training loss',
                              running_loss / 1000,
                              epoch * len(trainloader) + i)

    # 保存模型
    model_path = "./model/resnet.pth"
    torch.save(testNet.state_dict(), model_path)

    # 加载模型
    testNet = TestNet()
    testNet.load_state_dict(torch.load(model_path))

    # 测试
    correct, total = 0, 0
    with torch.no_grad():
        for data in trainloader:
            images, labels = data
            outputs = testNet(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    print('Accuracy of the network on the 10000 test images: %d %%' % (100 * correct / total))

    sys.exit(1)
